package edu.ivanov.apm.quant;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * User: Илья
 * Date: 06.10.13
 * Time: 18:49
 */
public class Solver {

	/** Краевые значения волновой функции */
	private static final double Y0 = 0;
	private static final double YN_1 = 0;

	/** Приближённые значения в узлах 1 и (n - 2) */
	private static final double Y1 = 1E-9;
	private static final double YN_2 = 1E-9;

	/** Начальный "отступ" от En_1 при поиске En */
	private static final double de = 1E-4;

	/** Приращение энергии при поиске отрезка, на котором f обращается в нуль */
	private static final double SEGMENT_SEARCH_STEP = 1E-4;

	/** Максимальное число итераций при поиске отрезка, на котором f обращается в нуль */
	private static final int SEGMENT_SEARCH_MAX_ITERATIONS = 10000;

	/** Максимальное количество итераций в методе дизотомииы */
	private static final int DICHOTOMY_MAX_ITERATIONS = 100000;

	/** масса частицы */
	private double m;
	/** половина длины потенциальной ямы */
	private double l;
	/** амплитуда потенциальной энергии */
	private double v0;
	/** число узлов разбиения равномерной сетки */
	private int n;
	/** минимальное значение потенциальной энергии */
	private double uMin;
	/** тончость пристрелки по значению энергии */
	private double eps;
	/** номер последнего квантовомеханического состояния, которое необходимо найти */
	private int i;

	/** шаг сетки */
	private double h;
	/** узел сшивки */
	private int r;
	/** константа h^2 / 12 для метода Нумерова */
	private double c;

	/** функция U(x) */
	private double u(double x) {
		return v0 * Math.sin(Math.PI * x / l);
	}

	/** функция q(E, x) */
	private double q(double e, double x) {
		return 2 * m * (e - u(x));
	}

	/**
	 * @param m Масса частицы.
	 */
	public void setM(double m) {
		this.m = m;
	}

	/**
	 * @param l L - модуль значений границ потенциальной ямы.
	 */
	public void setL(double l) {
		this.l = l;
	}

	/**
	 * @param v0 Амплитуда функции U(x).
	 */
	public void setV0(double v0) {
		this.v0 = v0;
	}

	/**
	 * @param n Число узлов сетки. Должно быть нечётным.
	 */
	public void setN(int n) {
		this.n = n;
	}

	/**
	 * @param uMin Минимум функции U(x) на отрезке [-L, L].
	 */
	public void setUMin(double uMin) {
		this.uMin = uMin;
	}

	/**
	 * @param eps Значение точности пристрелки.
	 */
	public void setEps(double eps) {
		this.eps = eps;
	}

	/**
	 * @param i Номер последнего квантовомеханического состояния, которое необходимо найти
	 */
	public void setI(int i) {
		this.i = i;
	}

	/**
	 * Получает волновые функции и соответствующие значения энергии,
	 * описывающие квантовомеханические состояния от 0 до заданного в файле настроек (параметр i).
	 *
	 * @return Решение.
	 * @throws SolverException Если решение не найдено по причине превышения допустимого числа итераций.
	 */
	public Solution[] solve() throws SolverException {
		if (m <= 0 ||
			l <= 0 ||
			v0 <= 0 ||
			n <= 1 ||
			n % 2 != 1 ||
			eps <= 0 ||
			i < 0)
			throw new SolverException("Значения параметров для решения уравнения некорректны");

		precalculate();

		List<Solution> solutions = new ArrayList<Solution>();

		double eEstimation = uMin;
		for (int k = 0; k <= i; k++) {
			double[] forward = new double[n];
			double[] backward = new double[n];

			double eStart = eEstimation + de;

			eEstimation = doSolve(forward, backward, eStart);
			double mean = findEnergyMean(forward);

			Solution solution = new Solution();
			solution.setL(l);
			solution.setForward(forward);
			solution.setBackward(backward);
			solution.setE(eEstimation);
			solution.setI(k);
			solution.setMean(mean);

			solutions.add(solution);
		}

		return solutions.toArray(new Solution[solutions.size()]);
	}

	private void precalculate() {
		h = 2 * l / (n - 1);
		r = n / 2 + n / 4;
		c = h * h / 12;
	}

	/**
	 * @param forward Массив для хранения решения при интегрировании вперёд.
	 * @param backward Массив для хранения решения при интегрировании назад.
	 * @param eStart Значение энергии, с которого начинается поиск собственного значения.
	 * @return Найденное собственное значение.
	 * @throws SolverException Если решение не найдено по причине превышения допустимого числа итераций.
	 */
	private double doSolve(double[] forward, double[] backward, double eStart) throws SolverException {
		// Ищем отрезок, на котором функция f меняет знак

		double fStart = f(forward, backward, eStart);

		double eFinish = eStart + SEGMENT_SEARCH_STEP;
		double fFinish = f(forward, backward, eFinish);

		int iteration = 0;
		while (fStart * fFinish > 0 && ++iteration <= SEGMENT_SEARCH_MAX_ITERATIONS) {
			eStart = eFinish;
			fStart = fFinish;

			eFinish += SEGMENT_SEARCH_STEP;
			fFinish = f(forward, backward, eFinish);
		}

		if (iteration > SEGMENT_SEARCH_MAX_ITERATIONS)
			throw new SolverException("Превышено допустимое число итераций при поиске отрезка, на котором f обращается в нуль");

		// Ищем методом дихотомии на отрезке [eStart, eFinish] значение энергии, при котором |f| <= eps

		double a = eStart; // Левая граница текущего отрезка
		double b = eFinish; // Правая граница текущего отрезка
		double m = (a + b) / 2; // Середина текущего отрезка

		// Соответствующие значения функций
		double fa = fStart;
		double fm = f(forward, backward, m);

		iteration = 0;
		while (Math.abs(fm) > eps && ++iteration <= DICHOTOMY_MAX_ITERATIONS)
			if (fa * fm < 0) {
				// Берём левый подотрезок
				b = m;

				m = (a + b) / 2;
				fm = f(forward, backward, m);
			} else {
				// Берём правый подотрезок
				a = m;
				fa = fm;

				m = (a + b) / 2;
				fm = f(forward, backward, m);
			}

		// Ситуация, которая не должна возникать при корректной реализации!
		if (iteration > DICHOTOMY_MAX_ITERATIONS)
			throw new SolverException("Превышено допустимое число итераций для метода дихотомии");

		return m;
	}

	private double f(double[] forward, double[] backward, double e) {
		numerov(forward, backward, e);
		scale(forward, backward);

		double value = derivative(forward, r) - derivative(backward, r);
		return value;
	}

	private void numerov(double[] forward, double[] backward, double e) {
		forward[0] = Y0;
		forward[1] = Y1;
		for (int i = 2; i < n; i++) {
			double xi_2 = - l + (i - 2) * h;
			double xi_1 = - l + (i - 1) * h;
			double xi = - l + i * h;
			forward[i] = (2 * (1 - 5 * c * q(e, xi_1)) * forward[i - 1] - (1 + c * q(e, xi_2)) * forward[i - 2]) / (1 + c * q(e, xi));
		}

		backward[n - 1] = YN_1;
		backward[n - 2] = YN_2;
		for (int i = n - 3; i >= 0; i--) {
			double xi = - l + i * h;
			double xi1 = - l + (i + 1) * h;
			double xi2 = - l + (i + 2) * h;
			backward[i] = (2 * (1 - 5 * c * q(e, xi1)) * backward[i + 1] - (1 + c * q(e, xi2)) * backward[i + 2]) / (1 + c * q(e, xi));
		}
	}

	private void scale(double[] forward, double[] backward) {
		// Ищём максимальный по модулю элемент forward и делим на него все элементы массива
		double max = Math.abs(forward[0]);
		for (int i = 1; i < n; i++) {
			double abs = Math.abs(forward[i]);
			if (abs > max)
				max = abs;
		}
		for (int i = 0; i < n; i++)
			forward[i] /= max;

		// Масштабируем backward так, чтобы в узле сшивки значения forward и backward совпадали.
		double coef = forward[r] / backward[r];
		for (int i = 0; i < n; i++)
			backward[i] *= coef;
	}

	private double derivative(double[] y, int i) {
		return (y[i - 2] - y[i + 2] + 8 * (y[i + 1] - y[i - 1])) / (12 * h);
	}

	private double secondDerivative(double[] y, int i) {
		double h2 = h * h;
		if (i == 0)
			return (y[i + 2] - 2 * y[i + 1] + y[i]) / h2; // Правая производная
		else if (i == y.length - 1)
			return (y[i] - 2 * y[i - 1] + y[i - 2]) / h2; // Левая производная
		else
			return (y[i + 1] - 2 * y[i] + y[i - 1]) / h2; // Центральная производная
	}

	/**
	 * Находит среднее квантовомеханическое значение полной энергии.
	 *
	 * @param psi Волновая функция.
	 * @return Среднее значение энергии.
	 */
	private double findEnergyMean(double[] psi) {
		psi = Arrays.copyOf(psi, psi.length);

		// Нормируем функцию psi
		double[] psi2 = new double[psi.length];
		for (int i = 0; i < psi2.length; i++)
			psi2[i] = psi[i] * psi[i];

		double norm = Math.sqrt(simpson(psi2));
		for (int i = 0; i < psi.length; i++)
			psi[i] /= norm;

		// Находим вектор для функции psi* H psi, где H -- оператор Гамильтона
		double[] y = new double[psi.length];
		for (int i = 0; i < y.length; i++) {
			double xi = - l + i * h;
			y[i] = psi[i] * (- secondDerivative(psi, i) / (2 * m) + u(xi) * psi[i]);
		}

		return simpson(y);
	}

	/**
	 * @param y Сеточная функция.
	 * @return Интеграл от функции на отрезке [-l, l], вычисленный по формуле Симпсона.
	 */
	private double simpson(double[] y) {
		double coefficient = h / 3;
		double integral = 0;
		int sup = n - 1;
		for (int i = 0; i <= sup; i++)
			if (i == 0 || i == sup)
				integral += y[i];
			else if (i % 2 != 0)
				integral += 4 * y[i];
			else
				integral += 2 * y[i];

		integral *= coefficient;
		return integral;
	}

	public static class SolverException extends Exception {

		public SolverException(String message) {
			super(message);
		}
	}
}
